ICLR｜基于3D几何信息的分子图表示学习

2021年9月29日，来自加拿大蒙特利尔大学魁北克人工智能研究所 (Mila) 的唐建等人在ICLR上发表文章，将3D几何信息整合到图自监督学习模型中，提出了图多视图预训练 (GraphMVP) 框架，为下游的分子性质预测任务提供了更精准的预测方法。

分子图表示学习是现代药物和材料发现中的一个基本问题。分子图通常由其2D拓扑结构进行建模，但最近研究发现，3D几何信息在预测分子功能中起着更为重要的作用。然而，真实场景中3D信息的缺乏严重阻碍了分子图对其几何图表示的学习。

为了应对这一挑战，本文提出了图多视图预训练(GraphMVP)框架，通过利用2D拓扑结构和3D几何视图之间的对应性和一致性来实现自监督学习 (SSL)。GraphMVP有效地学习了一个通过更丰富和更具鉴别力的3D几何视图进行增强的2D分子图编码器。本文也提供了理论上的说明来进一步证明GraphMVP的有效性。最后，综合实验表明，GraphMVP方法的性能优于现有的图SSL方法。代码可以在GitHub上获得：

https://github.com/chao1224/GraphMVP

介绍

由于3D结构编码具有的分子能量信息，分子性质可以通过3D几何结构来进行更好地预测。然而一般在在下游任务中，分子的立体化学结构往往是非常昂贵且稀缺的。为了解决这个问题，本文提出了图多视图预训练 (GraphMVP) 框架，在预训练2D分子编码器时利用了3D几何信息，之后在没有3D信息的情况下对下游任务进行微调。本文的学习范式是，在预先训练期间将3D分子结构的知识注入2D分子图编码器，使下游任务即使没有3D信息可用的情况下也可以受益于模型隐含的3D几何知识。

一般来说，图自监督学习 (SSL) 大致分为两类: 对比式SSL和生成式SSL，它们在监督信号的设计上有所不同。对比式图SSL在图间水平构造监督信号，并通过与其他图的对比来学习表示，而生成式图SSL侧重于在图内水平重建原始图。这两个SSL任务所获得的知识是相辅相成的，因此本文的GraphMVP框架将它们整合在一起，形成了一个更具有鉴别力的2D分子图表示。将本文的工作与现有方法区分开来的最重要的区别之一是，所有以前的方法仅仅关注于2D分子拓扑结构。然而，对于科学任务，如分子性质预测，3D几何的结合可以提供补充更全面的信息。为了填补这一空白，本文提出了利用GraphMVP的3D几何图自监督预训练模型。

本文的贡献包括：(1)是第一个将3D几何信息整合到图SSL中的模型；(2)提出了一个对比式SSL的和一个生成式的SSL任务用于预训练。之后，本文阐述了它们之间的差异，并通过实验验证了两者的结合可以使模型得到更好的效果；(3)提供了理论上的见解和案例研究来证明为什么增加3D几何信息是有益的；(4)在所有SSL方法中实现了SOTA的性能。

GraphMVP: 图多视图预训练

利用3D信息进行自监督学习 (SSL) 预训练，3D构象编码丰富的分子能量和空间结构的信息可以作为2D拓扑信息的补充。因此，在2D和3D视图中应用SSL方法将提供一个更好的2D分子表示，它隐含地将能量和几何信息的集合嵌入了分子表示中。

GraphMVP概述

图1: GraphMVP中预训练阶段的概述。黑色虚线圆圈表示子图掩码，本文在2D和3D图中对相同区域进行了掩码。分子的多个视图通过2D和3D GNN模型映射到表示空间，利用对比式和生成式代理任务，使用了GraphMVP执行SSL的预训练

通常，GraphMVP将每个分子2D拓扑和3D几何作为两个互补的视图。通过在这些视图之间执行SSL，模型期望通过3D构象来增强2D表示，这可以更好地反映分子的某些特性。

类似一般的SSL预训练pipeline，GraphMVP分为两个阶段：预训练之后微调。在预训练阶段，对辅助任务收集到的2D和3D分子结构数据进行SSL处理。在微调过程中，预先训练好的2D GNN模型随后在特定的下游任务上进行微调，其中只有2D分子图可用。

在SSL预训练阶段，本文设计了两个代理任务：一个对比式任务和一个生成式任务。本文推测并实证证明了这两个任务关注的是不同的方面的学习，这些方面可以归纳为以下两点。(1)从表示学习的角度看，对比式SSL利用数据间知识，生成式SSL利用数据内知识。对于对比式SSL一个关键的步骤是获得负视图对进行数据间的对比；而生成式SSL关注于每个数据点本身，在数据内水平重构关键特性。(2)从分布学习的角度来看，对比式SSL和生成式SSL分别从局部和全局的角度来学习数据分布。对比式SSL通过对比数据间成对距离来了解局部分布。因此，只要有足够多的数据点，局部对比运算就可以迭代地恢复数据分布。另一方面，生成式SSL直接学习全局数据密度函数。

因此，对比式的和生成式的SSL本质上是用不同的直觉和规则进行表示和分布式学习，本文期望两者的结合可以产生更好的分子表示。

2D与3D视图的对比式自监督学习

对比式自监督学习 (SSL) 的主要思想是首先从数据间的水平定义正视图对和负视图对，然后同时对齐正视图对和负视图对。对于每个分子，本文首先从2D和3D视图中提取分子表示。然后本文为对比式学习创建正负对：同一分子的2D-3D对被视为正对，否则被视为负对。最后，本文对正对和负对进行对比，pipeline如图1所示。

在对比式图SSL中，我们采用了基于能量的噪声对比估计模型 (EBM-NCE)。基于能量的噪声对比估计模型 (EBM-NCE) 是图对比SSL中广泛使用的一种替代方法。其目的与InfoNCE基本相同，即对正对和对比负对进行对齐，而主要区别在于对于负采样采用了二进制交叉熵和额外噪声分布，其中pn是噪声分布，σ是sigmoid函数：

2D与3D视图的生成式自监督学习

生成式SSL是另一个用于无监督预训练的经典方法。它的目的是通过对每个数据点进行自重构来学习一种有效的表示方法。具体到药物发现，本文对每个分子都有一个2D图和的一定数量的3D构象，本文的目标是学习一个鲁棒的2D/3D表示，可以在最大程度上恢复其3D/2D对应物。通过这样做，生成式 SSL可以强制2D/3D GNN对最关键的几何/拓扑信息进行编码，从而提高下游任务性能。

生成式模型有许多选择，包括变分自动编码器 (VAE)，生成式对抗网络 (GAN)，基于流的模型等。在 GraphMVP中，本文更倾向于 VAE 类方法，原因如下：(1)两个分子视图之间的映射是随机的：多个3D构象对应于相同的2D拓扑; (2)下游任务需要一个显式的2D图表示(即特征编码器); (3)结构化数据的解码器，例如图，通常是非常重要的设计，这使得它们成为次优选择。

因此，本文提出了一个轻量级类VAE的生成式SSL，并提出了一种新的代理损失名为变分表示重构 (VRR)，来将数据空间的重构转变为表示空间的重构：

它可以直接将隐码映射到3D表示空间，而不是将隐码解码到数据空间，这样就减小了从数据空间再变换到到表示空间带来的误差。

VRR实际上可以使MI最大化，而MI是连续双射函数的不变量。因此，如果编码函数满足这个条件，那么这个代理损失将是精确的。然而，我们发现GNN虽然不能满足这个条件，但能够提供相当鲁棒的性能，这从经验上证明了VRR的有效性。

多任务目标函数

如前所述，对比SSL和生成式SSL实质上是从不同的观点学习表示法。一个合理的推测是，结合两种SSL方法可以得到更好的总体性能：

现有的图SSL方法只关注于2D拓扑，这与GraphMVP是可以同时进行的：2D图SSL侧重于开发2D结构拓扑，GraphMVP利用了3D几何信息。因此，我们建议将2D SSL合并到GraphMVP中。由于2D图SSL主要有两大类：生成式和对比式，因此我们提出了两个模型变体GraphMVP-G和GraphMVP-C：

后续的实证结果也支持了GraphMVP-G和GraphMVP-C的有效性，因此，我们可以得出结论，现有的2D SSL可以对GraphMVP进行补充。

实验和结果

本文在相同的数据集上进行模型的预训练，然后对下游任务进行大范围的微调。本文从GEOM中随机选择了具有2D和3D结构的50k个合格分子进行预训练。由于构象集合可以更好地反映分子性质，因此本文对每个分子取C个构象。对于下游任务，本文首先坚持与当今主流的SSL工作采取相同的设置，探索8个二元分子性质预测任务，这些任务都具有低数据。然后，本文从各种低数据领域探讨了6个回归任务使分析更加全面。

表1：分子性质预测任务的结果。对于每个下游任务，本文报告了3个支架分裂种子的 ROC-AUC 平均值(和标准差)。对于 GraphMVP，本文设置 M = 0.15和 C = 5。最佳和次佳结果分别标记为粗体和粗体。

对于分子性质预测的主要结果

本文对10个最公认的SSL baseline和随机初始化进行了综合的比较。在预训练阶段，本文基于GEOM将所有SSL方法应用于同一个数据集上。在微调阶段，本文与8个低数据分子性质预测任务采用了相同的设置。对于GraphMVP模型有两个关键因素: 掩蔽率 (M) 和每个分子的构象数 (C)，本文默认设置 M = 0.15和 C = 5。

表1列出了8个分子性质预测任务的主要结果。可以观察到GraphMVP的性能明显优于随机初始化方法，平均性能大大优于现有的SSL方法。此外，GraphMVP-G和GraphMVP-C都可以改善模型的性能，实验结果可以得到这样的结论：3D几何图形与2D拓扑是互补的。GraphMVP利用3D几何和2D拓扑之间的信息，2D SSL作为正则化项提取更多的2D拓扑信息；它们提取信息的不同视角，实际上是相互补充的。

更广泛的下游任务

目前讨论的8个二分类下游任务已广泛应用于分子的图SSL研究，但还有更多任务3D构象也可以提供帮助。在这里，我们测试了4个额外的回归性质预测任务和2个药物-靶点亲和力任务。

药物-靶点亲和力 (DTA) 是药物发现中的一项关键任务，它对分子药物和靶蛋白进行建模，目的是预测他们的亲和力分数。最近的一项工作是使用2D GNN对分子药物进行建模，使用卷积神经网络(CNN)对和靶蛋白 (作为一个氨基酸序列) 进行建模。我们使用GraphMVP对2D GNN进行预训练来与这项工作采取同样的设置。如表2所示，一致的性能增益验证了我们提出的GraphMVP模型的有效性。

表2: 四个分子特性预测任务(回归)和两个 DTA 任务(回归)的结果。我们报告了3个种子的平均均方根误差(RMSE)和3个种子的平均均方根误差(MSE)。对于 GraphMVP，我们设置 M = 0.15和 C = 5。每个任务的表现最佳的模型用粗体标出。我们在这里省略了std，因为它们非常小而且难以区分。

案例研究

本文研究了在利用2D拓扑来解决具有困难，但使用3D几何图形来解决非常简单的案例中 (如图2所示)，GraphMVP是如何起作用的。因此，本文设计了两个案例来验证GraphMVP如何将知识从3D几何转化为2D表示。

第一个案例是3D直径预测。对于分子，通常2D直径越长，3D直径越大 (最大的原子成对距离 l2)。然而，这并不总是成立，本文尝试使用2D图预测3D直径。第二个案例研究是远距离供体-受体检测。分子具有一种特殊的几何结构，称为供体-受体键，本文希望利用2D分子图来检测这种特殊的结构。经验证GraphMVP在这两个案例研究中都对性能带来了改进。

图2: 本文选择的分子，其性质可以很容易地解决通过3D而不是2D来发现。随机初始化的2D GNN 精度分别为38.9 ± 0.8和77.9 ± 1.1。预训练的 GraphMVP 得分分别为42.3 ± 1.3和81.5 ± 0.4，优于所有先例。本文绘制了随机初始化预测失败但 GraphMVP 预测正确的情况

结论和未来工作

在这项工作中，我们提供了一个非常通用的框架，称为 GraphMVP。

从领域的角度来看，GraphMVP是第一个结合三维信息来增强二维图形表示学习的算法，并且在建模中考虑了随机性，能够利用三维一致性的优势。

在技术创新方面，GraphMVP在引入2个SSL任务时带来了以下见解:(1) GraphMVP提出了EBM-NCE和VRR，分别使用EBM和变分分布对条件分布进行建模。(2)EBM-NCE与JSE相似，虽然我们对理论直觉的研究方向不同，但EBM-NCE在这一领域开辟了另一个有前途的研究场景。VRR作为一种生成SSL方法，能够缓解分子生成中的潜在问题。最后，GraphMVP将对比式SSL (InfoNCE 或 EBM-NCE)和生成式SSL (VRR) 结合起来实现目标函数。实证结果 (对14个下游数据集的可靠性改进) 和理论分析都可以有力地支持上述领域和技术上的贡献。

我们想强调的是，GraphMVP是模型无关的，并且有可能扩展到许多其他低数据量的应用中。这推动了未来广泛的探索方向，包括但不限于: (1)更强大的二维和三维分子表示方法。(2)小分子以外的不同应用领域，如蛋白质等大分子。

参考资料

Pre-training Molecular Graph Representation with 3D Geometry

https://openreview.net/forum?id=xQUe1pOKPam

--------- End ---------

感兴趣的读者，可以添加小邦微信（zhiyaobang2020）加入读者实名讨论微信群。添加时请主动注明姓名-企业-职位/岗位 或  姓名-学校-职务/研究方向。

- 历史文章推荐 -

【分子表征】

【AI药物设计】